package com.wave.wavesystem.ai.config;

import com.wave.common.util.Func;
import com.wave.wavesystem.ai.enums.SSEMsgType;
import lombok.extern.slf4j.Slf4j;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

@Slf4j
@Component
public class SSEConfig {

    private static final Map<String, SseEmitter> sseClientMap = new ConcurrentHashMap<>();

    private static final Map<String, Integer> reConnectCountMap = new ConcurrentHashMap<>();

    /**
     * 链接sse 服务
     *
     * @param conversationId 会话id
     */
    public static SseEmitter connect(String conversationId) {
        if (sseClientMap.containsKey(conversationId)) {
            return sseClientMap.get(conversationId);
        }
        SseEmitter sseEmitter = new SseEmitter();
        sseEmitter.onTimeout(timeoutCallback(conversationId));
        sseEmitter.onCompletion(completionCallback(conversationId));
        sseEmitter.onError(errorCallback(conversationId));
        sseClientMap.put(conversationId, sseEmitter);
        log.info("SSE connect success. conversationId={}", conversationId);
        return sseEmitter;

    }

    public static Runnable timeoutCallback(String conversationId) {
        return () -> {
            // 移出链接
            remove(conversationId);
            log.info("sse链接超时，conversationId：{}", conversationId);
        };
    }

    public static Runnable completionCallback(String conversationId) {
        return () -> {
            // 移出链接
            remove(conversationId);
            log.info("sse链接完成，conversationId：{}", conversationId);
        };
    }

    public static void sendMessage(String conversationId,
                                   String message, SSEMsgType type) {
        if (Func.isEmpty(sseClientMap)) {
            return;
        }
        if (sseClientMap.containsKey(conversationId)) {
            SseEmitter sseEmitter = sseClientMap.get(conversationId);
            sendSseEmitterMessage(sseEmitter, conversationId, message, type);
        }
    }

    private static void sendSseEmitterMessage(SseEmitter sseEmitter, String conversationId,
                                              String message, SSEMsgType type) {
        SseEmitter.SseEventBuilder sseEventBuilder = SseEmitter.event()
                .id(conversationId)
                .data(message)
                .name(type.getType());
        try {
            sseEmitter.send(sseEventBuilder);
        } catch (IOException e) {
            remove(conversationId);
            log.info("sse 异常：{}", e.getMessage());
        }

    }

    public static void complete(String conversationId) {
        if (sseClientMap.containsKey(conversationId)) {
            SseEmitter sseEmitter = sseClientMap.get(conversationId);
            sseEmitter.complete();
        }
        log.info("数据推送完成！：{}", conversationId);

    }


    public static Consumer<Throwable> errorCallback(String conversationId) {
        return throwable -> {
            remove(conversationId);
            log.info("sse链接异常，conversationId：{}", conversationId);
        };
    }

    private static void remove(String conversationId) {
        // 删除用户
        sseClientMap.remove(conversationId);
        log.info("sse链接移出，conversationId：{}", conversationId);
    }


    @Scheduled(fixedRate = 15_000)
    public void ping() {
        if (Func.isEmpty(sseClientMap)) {
            return;
        }
        for (Map.Entry<String, SseEmitter> en : sseClientMap.entrySet()) {
            String cid = en.getKey();
            SseEmitter emitter = en.getValue();
            try {
                // 发一条 SSE 注释行，浏览器不会触发 onmessage
                emitter.send(SseEmitter.event().comment("ping"));
                log.info("SSE 心跳检测中 {} ", cid);
            } catch (Exception e) {
                if (reConnectCountMap.containsKey(cid)) {
                    int count = reConnectCountMap.get(cid) + 1;
                    if (count >= 3) {
                        remove(cid);
                        log.info("SSE 重试次数大于3次 移出 {} removed", cid);
                    } else {
                        reConnectCountMap.put(cid, count);
                        log.info("SSE 重试次数累加 {} ，{}次", cid, count);
                    }
                    return;
                }
                reConnectCountMap.put(cid, 1);
                log.info("SSE 重试次数累加 {} ", cid);
            }
        }
    }

}
